/*
 * @Project Name: web-official-cms
 * @File Name: TokenBucketUtil.java
 * @Package Name: com.hhly.sns.cms.util
 * @Date: 2017年9月18日上午8:41:39
 * @Creator: xuyuji-374
 * @line------------------------------
 * @修改人: 
 * @修改时间: 
 * @修改内容: 
 */

package tool;

import java.util.concurrent.ConcurrentHashMap;

/**
 * 令牌桶工具类
 * 
 * @author xuyuji-374
 * @date 2017年9月18日上午8:41:39
 * @see
 */
public class TokenBucketUtil {

	private static final int BUCKET_CAPACITY_DEFAULT = 10;
	private static final int RECOVERY_SPEED_DEFAULT = 1000 * 60 / 10;
	private static final int CACHE_CAPACITY_DEFAULT = 1065536;

	private int bucketCapacity; // 令牌桶容量
	private int recoverySpeed; // 令牌恢复时间（毫秒）
	private int cacheCapacity; // 缓存令牌桶的哈希表容量
	private ConcurrentHashMap<String, TokenBucket> tokenBucketCache;

	public TokenBucketUtil() {
		bucketCapacity = BUCKET_CAPACITY_DEFAULT;
		recoverySpeed = RECOVERY_SPEED_DEFAULT;
		cacheCapacity = CACHE_CAPACITY_DEFAULT;
		tokenBucketCache = new ConcurrentHashMap<>(cacheCapacity);
	}

	public TokenBucketUtil(int bucketCapacity, int recoverySpeed, int cacheCapacity) {
		if (bucketCapacity > 0) {
			this.bucketCapacity = bucketCapacity;
		} else {
			this.bucketCapacity = BUCKET_CAPACITY_DEFAULT;
		}

		if (recoverySpeed > 0) {
			this.recoverySpeed = recoverySpeed;
		} else {
			this.recoverySpeed = RECOVERY_SPEED_DEFAULT;
		}

		if (cacheCapacity > 0) {
			this.cacheCapacity = cacheCapacity;
		} else {
			this.cacheCapacity = CACHE_CAPACITY_DEFAULT;
		}
		tokenBucketCache = new ConcurrentHashMap<>(this.cacheCapacity);
	}

	public int getBucketCapacity() {
		return bucketCapacity;
	}

	public void setBucketCapacity(int bucketCapacity) {
		if (bucketCapacity > 0) {
			this.bucketCapacity = bucketCapacity;
		}
	}

	public int getRecoverySpeed() {
		return recoverySpeed;
	}

	public void setRecoverySpeed(int recoverySpeed) {
		if (recoverySpeed > 0) {
			this.recoverySpeed = recoverySpeed;
		}
	}

	public int getCacheCapacity() {
		return cacheCapacity;
	}

	public void setCacheCapacity(int cacheCapacity) {
		if (cacheCapacity > 0) {
			this.cacheCapacity = cacheCapacity;
			this.tokenBucketCache = new ConcurrentHashMap<>(cacheCapacity);
		}
	}

	public boolean visit(String key) {
		if (tokenBucketCache.size() > cacheCapacity) {
			tokenBucketCache = new ConcurrentHashMap<>(cacheCapacity);
		}
		long now = System.currentTimeMillis();
		TokenBucket tokenBucket = tokenBucketCache.get(key);
		if (tokenBucket == null) {
			tokenBucket = new TokenBucket(now, bucketCapacity);
			tokenBucketCache.put(key, tokenBucket);
		}
		int nowTokenNum = (int) (tokenBucket.getTokenNum() + (now - tokenBucket.getLastTime()) / recoverySpeed);
		if (nowTokenNum > bucketCapacity) {
			nowTokenNum = bucketCapacity;
		}
		if (nowTokenNum > 0) {
			tokenBucket.setLastTime(now);
			tokenBucket.setTokenNum(nowTokenNum - 1);
			return true;
		} else {
			return false;
		}
	}

	private class TokenBucket {
		private long lastTime;
		private int tokenNum;

		public TokenBucket(long lastTime, int tokenNum) {
			this.lastTime = lastTime;
			this.tokenNum = tokenNum;
		}

		public long getLastTime() {
			return lastTime;
		}

		public void setLastTime(long lastTime) {
			this.lastTime = lastTime;
		}

		public int getTokenNum() {
			return tokenNum;
		}

		public void setTokenNum(int tokenNum) {
			this.tokenNum = tokenNum;
		}
	}

	public static void main(String[] args) throws InterruptedException {
		TokenBucketUtil tokenBucketUtil = new TokenBucketUtil(3, 2000, 1024);
		while (true) {
			System.out.println("xuyuji visit:" + tokenBucketUtil.visit("xuyuji"));
			Thread.sleep(1000L);
		}
	}
}
